import os
import subprocess

def run_full_pipeline(fits_dir, cand_output_dir, filtered_txt_path, plot_output_dir, save_npz=False):
    os.environ['FITS_DIR'] = fits_dir
    os.environ['CAND_DIR'] = cand_output_dir
    os.environ['TXT_PATH'] = filtered_txt_path
    os.environ['PLOT_DIR'] = plot_output_dir

    print(f"\n======Running pipeline for: {fits_dir}======")
    print(f"===> Cand Dir:    {cand_output_dir}")
    print(f"===> Filtered:    {filtered_txt_path}")
    print(f"===> Plot Output: {plot_output_dir}")

    try:
        # Step 1: Run Heimdall
        subprocess.run(["bash", "run_heimdall_batch.sh", fits_dir, cand_output_dir], check=True)

        # Step 2: Filter candidates
        subprocess.run(["python", "filter_cands_simple.py", cand_output_dir, filtered_txt_path, fits_dir], check=True)

        # Step 3: Visualization
        args = ["python", "run_frb_search.py", fits_dir, filtered_txt_path, plot_output_dir]
        if save_npz:
            args.append("--save_npz")
        subprocess.run(args, check=True)

        print("======Completed pipeline for:", fits_dir, "======")
    except subprocess.CalledProcessError as e:
        print(f"Error occurred while processing {fits_dir}. Skipping this directory...")
        return  # Skip to the next directory

def should_skip_directory(date_dir, skip_list):
    """检查当前目录是否在跳过列表中"""
    return date_dir in skip_list

if __name__ == "__main__":
    # 根路径：硬盘中数据路径
    base_data_dir = "/media/user/niuchenhui18/PT2022_0034/FRB190520"
    project_name = os.path.basename(os.path.dirname(base_data_dir))  # => 项目名称

    # 工作目录：结果保存路径
    workspace_root = os.path.expanduser(f"~/workspace/FRB_search/FRB190520/{project_name}")

    # 指定需要跳过的目录列表（直接在程序中指定）
    skip_dirs = ["20220805"]  # 这里列出你要跳过的目录

    # 遍历日期文件夹
    for date_dir in sorted(os.listdir(base_data_dir)):
        full_fits_dir = os.path.join(base_data_dir, date_dir)
        if not os.path.isdir(full_fits_dir):
            continue  # 过滤掉非目录项

        # 检查是否跳过当前目录
        if should_skip_directory(date_dir, skip_dirs):
            print(f"Skipping directory: {date_dir}")
            continue

        # 设置输出路径（按日期分类）
        output_root = os.path.join(workspace_root, date_dir)
        cand_output_dir = os.path.join(output_root, "cand")
        filtered_txt_path = os.path.join(output_root, "filtered.txt")
        plot_output_dir = os.path.join(output_root, "plot")

        # 创建必要的目录
        os.makedirs(cand_output_dir, exist_ok=True)
        os.makedirs(plot_output_dir, exist_ok=True)

        # 执行主处理流程
        run_full_pipeline(
            fits_dir=full_fits_dir,
            cand_output_dir=cand_output_dir,
            filtered_txt_path=filtered_txt_path,
            plot_output_dir=plot_output_dir,
            save_npz=False
        )